import os
import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np

K = np.array(
    [
        [
            0.19611235,
            0.36949325,
            0.8252357,
            0.60269207,
            0.853726,
            0.14729041,
            0.93213016,
            0.4912518,
            0.4421994,
            0.94502634,
        ],
        [
            0.8472717,
            0.08033702,
            0.7874011,
            0.37478322,
            0.7154022,
            0.36539742,
            0.42748713,
            0.46491373,
            0.12844378,
            0.4038347,
        ],
        [
            0.7179675,
            0.5542868,
            0.12964171,
            0.310674,
            0.28132373,
            0.4555802,
            0.48048112,
            0.19307618,
            0.3653659,
            0.3019691,
        ],
        [
            0.5715858,
            0.45215115,
            0.50888586,
            0.44008783,
            0.6940242,
            0.09332542,
            0.7861716,
            0.07335378,
            0.93303907,
            0.9602915,
        ],
        [
            0.37165388,
            0.41009933,
            0.34311947,
            0.86977965,
            0.03910041,
            0.23821034,
            0.7766873,
            0.03240556,
            0.7266301,
            0.8464755,
        ],
        [
            0.7198575,
            0.22694403,
            0.04758935,
            0.7079051,
            0.39066452,
            0.54978186,
            0.46865922,
            0.67030114,
            0.15897818,
            0.8474187,
        ],
        [
            0.8557079,
            0.31147653,
            0.9369019,
            0.668592,
            0.52999514,
            0.62273526,
            0.5848188,
            0.6915461,
            0.86041534,
            0.02546168,
        ],
        [
            0.7917908,
            0.91307646,
            0.78238267,
            0.3207483,
            0.8921884,
            0.450409,
            0.1324447,
            0.67176586,
            0.89350206,
            0.9026417,
        ],
        [
            0.65101486,
            0.2760342,
            0.3508995,
            0.5330928,
            0.8670503,
            0.0992681,
            0.42444643,
            0.45056775,
            0.75550824,
            0.53870803,
        ],
        [
            0.29493734,
            0.7684385,
            0.34008777,
            0.71705914,
            0.99242055,
            0.8508037,
            0.85793114,
            0.588454,
            0.9230404,
            0.16057819,
        ],
    ]
)


# 定义神经网络模型
class DeepNN(nn.Module):
    def __init__(self, input_dim):
        """
        初始化DeepNN类。

        参数:
        input_dim (int): 输入数据的维度。

        返回:
        无

        说明:
        该构造函数创建了一个深度神经网络模型，包括输入层、多层隐藏层和输出层。
        每层使用线性单元，神经元数量逐层递减，最后输出层的神经元数量与输入维度相同。
        """
        # 调用父类的构造函数
        super(DeepNN, self).__init__()

        # 定义第一层线性单元，将输入维度映射到128维
        self.fc1 = nn.Linear(input_dim, 128)

        # 定义第二层线性单元，将128维映射到64维
        self.fc2 = nn.Linear(128, 64)

        # 定义第三层线性单元，将64维映射到32维
        self.fc3 = nn.Linear(64, 32)

        # 定义第四层线性单元，将32维映射到16维
        self.fc4 = nn.Linear(32, 16)

        # 定义输出层线性单元，将16维映射回输入维度
        self.fc5 = nn.Linear(16, input_dim)

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        x = torch.relu(self.fc3(x))
        x = torch.relu(self.fc4(x))
        x = self.fc5(x)
        return x


# 生成随机的K矩阵和u向量
def generate_data(num_samples, input_dim):
    K = np.random.rand(input_dim, input_dim).astype(np.float32)  # 生成固定的K
    u = np.random.rand(num_samples, input_dim).astype(np.float32)
    F = np.dot(u, K)  # 计算 F = u * K
    return u, F, K


def main():
    # 超参数设置
    input_dim = 10
    num_samples = 100000
    epochs = 100
    learning_rate = 0.001

    pth_file = "model.pth"

    model = DeepNN(input_dim)
    # 生成数据
    u, F, K = generate_data(num_samples, input_dim)

    if os.path.exists(pth_file):
        # 加载模型
        model.load_state_dict(torch.load(pth_file, weights_only=True))
    else:

        # 转换为PyTorch张量
        u_tensor = torch.tensor(u)
        F_tensor = torch.tensor(F)

        # 定义模型、损失函数和优化器
        criterion = nn.MSELoss()
        optimizer = optim.Adam(model.parameters(), lr=learning_rate)

        # 训练模型
        for epoch in range(epochs):
            model.train()

            optimizer.zero_grad()

            outputs = model(u_tensor)
            loss = criterion(outputs, F_tensor)

            loss.backward()
            optimizer.step()

            if (epoch + 1) % 10 == 0:
                print(f"Epoch [{epoch + 1}/{epochs}], Loss: {loss.item():.4f}")

        # 保存模型参数
        torch.save(model.state_dict(), pth_file)

    # 生成随机u向量进行预测
    u_test = np.random.rand(1, input_dim).astype(np.float32)
    u_test_tensor = torch.tensor(u_test)

    # 模型预测
    model.eval()
    with torch.no_grad():
        F_pred_tensor = model(u_test_tensor)
        F_pred = F_pred_tensor.numpy()

    print("测试u向量:", u_test)
    print("预测的F向量:", F_pred)
    print("实际的F向量:", np.dot(u_test, K))


if __name__ == "__main__":
    main()
